from lark import Lark, Transformer

from PBLTL.language import AP, Bool, And, Or, Not, Imply, Iff, U, X, F, G, R, P, Interval

pbltl_grammar = r"""

?path_formula : prob_expr
| "X" path_formula -> next_expr
| "F" interval path_formula -> final_expr
| "G" interval path_formula -> global_expr
| path_formula "R" interval path_formula -> release_expr
| path_formula "U" interval path_formula -> until_expr
| path_formula "&" path_formula -> and_expr
| path_formula "|" path_formula -> or_expr
| "~" path_formula -> not_expr
| path_formula "->" path_formula -> imply_expr
| path_formula "<->" path_formula -> iff_expr
| "(" path_formula ")" -> sub_expr
| "true" -> true_expr
| "false" -> false_expr
| ap

interval: "[" number "," number_or_inf "]" -> interval_lcrc
        | "[" number "," number_or_inf ")" -> interval_lcro
        | "(" number "," number_or_inf "]" -> interval_lorc
        | "(" number "," number_or_inf ")" -> interval_loro

number_or_inf: number -> number
             | "inf"  -> inf

prob_expr: "P" "=?" "{" path_formula "}"            -> prob_expr_cal
         | "P" "<" number "{" path_formula "}"      -> prob_expr_less
         | "P" "<=" number "{" path_formula "}"     -> prob_expr_less_equal
         | "P" ">" number "{" path_formula "}"      -> prob_expr_greater
         | "P" ">=" number "{" path_formula "}"     -> prob_expr_greater_equal
         
ap: string"."string"="string            -> ap_equal
  | string"."string"="number            -> ap_equal
  | string"."string"<"number            -> ap_less
  | string"."string"<="number           -> ap_less_equal
  | string"."string">"number            -> ap_greater
  | string"."string">="number           -> ap_greater_equal

string: /[a-zA-Z_][a-zA-Z0-9_]*/ 

number: NUMBER -> number
 

%import common.NUMBER
%import common.ESCAPED_STRING
%import common.WS
%ignore WS

"""

class PBLTLTransformer(Transformer):
    def true_expr(self, _):
        return Bool(True)

    def false_expr(self, _):
        return Bool(False)

    def string(self, args):
        (s,) = args
        return s[:]

    def inf(self, args):
        return float("inf")

    def number(self, args):
        (n,) = args
        return float(n)

    def interval_lcrc(self, args):
        (n1, n2) = args
        return Interval("[", float(n1), float(n2), "]")

    def interval_lcro(self, args):
        (n1, n2) = args
        return Interval("[", float(n1), float(n2), ")")

    def interval_lorc(self, args):
        (n1, n2) = args
        return Interval("(", float(n1), float(n2), "]")

    def interval_loro(self, args):
        (n1, n2) = args
        return Interval("(", float(n1), float(n2), ")")

    def ap_equal(self, args):
        return AP(args[0], args[1], "=", args[2])

    def ap_less(self, args):
        return AP(args[0], args[1], "<", args[2])

    def ap_less_equal(self, args):
        return AP(args[0], args[1], "<=", args[2])

    def ap_greater(self, args):
        return AP(args[0], args[1], ">", args[2])

    def ap_greater_equal(self, args):
        return AP(args[0], args[1], ">=", args[2])

    def not_expr(self, args):
        return Not(args[0])
    
    def and_expr(self, args):
        return And(args[0], args[1])

    def or_expr(self, args):
        return Or(args[0], args[1])

    def imply_expr(self, args):
        return Imply(args[0], args[1])

    def iff_expr(self, args):
        return Iff(args[0], args[1])

    def sub_expr(self, args):
        return args[0]

    def final_expr(self, args):
        return F(args[0], args[1])

    def global_expr(self, args):
        return G(args[0], args[1])

    def release_expr(self, args):
        return R(args[1], args[0], args[2])

    def until_expr(self, args):
        return U(args[1], args[0], args[2])

    def prob_expr_cal(self, args):
        return P(args[0], "=?")

    def prob_expr_less(self, args):
        return P(args[1], "<", args[0])

    def prob_expr_less_equal(self, args):
        return P(args[1], "<=", args[0])

    def prob_expr_greater(self, args):
        return P(args[1], ">", args[0])

    def prob_expr_greater_equal(self, args):
        return P(args[1], ">=", args[0])

    def next_expr(self, args):
        return X(args[0])

parser = Lark(pbltl_grammar, start="path_formula")


def parse(s: str):
    try:
        tree = parser.parse(s)
        result = PBLTLTransformer().transform(tree)
    except Exception as e:
        print(f"parse error in string: {s}, error: {e}")
        result = None
    return result

if __name__ == "__main__":
    sl = [
            "true",
            "false",
            "n1.state=work",
            "~n1.state=work",
            "P=?{n1.state=work U(0,0.2] n1.state=fail}",
            "P=?{F[0,0.2] n1.state=fail} ",
            "P=?{X n1.state=fail}",
            "P=?{G[0,0.2] n1.state=fail}",
            "P=?{n1.state=work R[0,0.2] s1.state=fail}",
            "P=?{n1.state=fail -> s1.state=fail}",
            "P=?{n1.state=work}",
            "P<0.2{n1.state=work}",
            "P<=0.2{n1.state=work}",
            "P>0.2{n1.state=work}",
            "P>=0.2{n1.state=work}",
    ]


    rl = [parse(s) for s in sl ]

    print(rl)
